"""
均线斜率峰值模块
负责计算和绘制移动平均线斜率峰值
"""

import pandas as pd
import numpy as np
import pyqtgraph as pg
from typing import List, Tuple, Dict, Optional
from PyQt5 import QtCore, QtGui
from ..BaseChart import BaseChart
from ..config import config, flags, logger

from collections import deque

class MovingAverageSlopePeaksCalculator:
    """移动平均线计算器"""
    
    def __init__(self, data: pd.DataFrame):
        """
        初始化移动平均线计算器
        
        Args:
            data: 包含股票数据的DataFrame
        """
        self.data = data.copy()
        self.ma_data = {}  # 存储计算好的均线数据
        self.ma_slope_peaks_data = {}  # 存储计算好的均线斜率峰值数据
        self.crest_trough_dit = {}  # [] 存储峰顶和谷底列表

    def calculate_ma_slope_peaks(self, periods: List[int], price_column: str = 'close') -> Dict[int, List[Tuple[int, int, int]]]:
        """
        计算指定周期的移动平均线斜率峰值
        
        Args:
            periods: 均线周期列表，如 [5, 10, 20, 30]
            price_column: 价格列名，默认为'close'
            
        Returns:
            包含各周期均线的字典
        """
        # ma_results = {}
        price = self.data[price_column]
        crest_trough_results = {}
        for period in periods:
            if period <= 0:
                continue
            
            ma_slope_peaks = []
            crest_trough_list = []
            # 计算移动平均线
            ma_series = self.data[price_column].rolling(window=period, min_periods=1).mean()
            tmp_ma = np.concatenate(([0], ma_series[:-1]))
            ma_slope_peaks = ma_series - tmp_ma
            # ma_slope_peaks.append(ma_series.diff().rolling(window=period, min_periods=1).mean())
            # ma_results[period] = ma_series
            # 存储到实例变量中
            self.ma_data[period] = ma_series
            self.ma_slope_peaks_data[period] = ma_slope_peaks
            l = len(ma_slope_peaks)
            max_idx_res, min_idx_res = self.sliding_window_max_min_index(price, period)
            for i, value in enumerate(ma_slope_peaks):
                if not pd.isna(value):
                    if i < l - 1 and not pd.isna(ma_slope_peaks[i + 1]):
                        t = ma_slope_peaks[i + 1]
                        if value > 0 and t < 0:  # 峰顶
                            ma_slope_peaks[i] = value
                            crest_trough_list.append((1, i, max_idx_res[i]))

                        elif value < 0 and t > 0: # 谷底
                            ma_slope_peaks[i] = value
                            crest_trough_list.append((-1, i, min_idx_res[i]))
            self.crest_trough_dit[period] = crest_trough_list
            crest_trough_results[period] = crest_trough_list
        # print(f"crest_trough_list: {crest_trough_results}")
        return crest_trough_results
    
    def sliding_window_max_min_index(self, arr, n):
        if n <= 0:
            return [], []
        
        max_idx_res = []
        min_idx_res = []
        dq_max = deque()  # 单调递减队列，队首索引最大
        dq_min = deque()  # 单调递增队列，队首索引最小

        for i, val in enumerate(arr):
            # ---------- 移除不在窗口内的索引 ----------
            while dq_max and dq_max[0] < i - n + 1:
                dq_max.popleft()
            while dq_min and dq_min[0] < i - n + 1:
                dq_min.popleft()
            
            # ---------- 维护单调性 ----------
            while dq_max and arr[dq_max[-1]] <= val:
                dq_max.pop()
            while dq_min and arr[dq_min[-1]] >= val:
                dq_min.pop()
            
            dq_max.append(i)
            dq_min.append(i)

            # ---------- 保存当前窗口最大值/最小值的索引 ----------
            max_idx_res.append(dq_max[0])
            min_idx_res.append(dq_min[0])
        
        return max_idx_res, min_idx_res
    
    def get_ma_data(self, period: int) -> Optional[pd.Series]:
        """获取指定周期的均线数据"""
        return self.ma_data.get(period)
    
    def get_ma_slope_peaks_data(self, period: int) -> Optional[pd.Series]:
        """获取指定周期的均线斜率峰值数据"""
        return self.ma_slope_peaks_data.get(period)
    
    def get_ma_range(self, periods: List[int]) -> Tuple[float, float]:
        """
        获取均线的价格范围
        
        Args:
            periods: 均线周期列表
            
        Returns:
            (最小价格, 最大价格)
        """
        min_price = float('inf')
        max_price = float('-inf')
        
        for period in periods:
            if period in self.ma_data:
                ma_series = self.ma_data[period]
                min_price = min(min_price, ma_series.min())
                max_price = max(max_price, ma_series.max())
        
        return min_price, max_price

class MovingAverageSlopePeaksItem(pg.GraphicsObject):
    """移动平均线绘制项"""
    
    def __init__(self, data: List[Tuple[int, float, float, float, float]], crest_trough_data: List[Tuple[int, int, int]], color: str = 'white', vb: pg.ViewBox = None, width: int = 1):
        """
        初始化移动平均线绘制项
        
        Args:
            data: K线数据，格式为 [(索引, 开盘, 收盘, 最低, 最高), ...]
            crest_trough_data: 波峰波谷坐标数据
            color: 线条颜色
            width: 线条宽度
        """
        pg.GraphicsObject.__init__(self)
        self.data = data
        self.crest_trough_data = crest_trough_data
        self.color = color
        self.width = width
        self.vb = vb
        self.generatePicture()
    
    def generatePicture(self):
        """生成拐点标志图形"""
        self.picture = QtGui.QPicture()
        p = QtGui.QPainter(self.picture)
        
        # 设置线条样式 - 使用更安全的颜色设置
        try:
            pen_1 = pg.mkPen('r', width=self.width)
            pen_2 = pg.mkPen('g', width=self.width)
        except:
            # 如果颜色解析失败，使用默认颜色
            pen_1 = pg.mkPen('white', width=self.width)
            pen_2 = pg.mkPen('white', width=self.width)
        # p.setPen(pen_1)

        # a = QtCore.QPointF(5, 5)
        # b = QtCore.QPointF(10, 10)

        # p.drawLine(a, b)
        # draw_equilateral_triangle(p, 5, 5, 3, self.vb, pen_1)
        # logger.debug(f"设置线条样式: {pen_1}")
        # 绘制均线斜率峰值
        # fh = 0 # 1 + ; -1 - ;
        high_price = self.data["high"]
        low_price = self.data["low"]
        for i, value in enumerate(self.crest_trough_data):
            # print(f"value: {value}")
            if value[0] == 1:
                p.setPen(pen_1)
                p.drawLine(QtCore.QPointF(value[2] - config.CANDLESTICK_WIDTH, high_price[value[2]]), QtCore.QPointF(value[2] + config.CANDLESTICK_WIDTH, high_price[value[2]]))
            elif value[0] == -1:
                p.setPen(pen_2)
                p.drawLine(QtCore.QPointF(value[2] - config.CANDLESTICK_WIDTH, low_price[value[2]]), QtCore.QPointF(value[2] + config.CANDLESTICK_WIDTH, low_price[value[2]]))
        # l = len(self.crest_trough_data)
        # for i, value in enumerate(self.crest_trough_data):
        #     if not pd.isna(value):
        #         if i < l - 1 and not pd.isna(self.crest_trough_data[i + 1]):
        #             t = self.crest_trough_data[i + 1]
        #             if value > 0 and t < 0:
        #                 p.setPen(pen_1)
        #                 p.drawLine(QtCore.QPointF(i - config.CANDLESTICK_WIDTH, high_price[i]), QtCore.QPointF(i + config.CANDLESTICK_WIDTH, high_price[i]))
        #             elif value < 0 and t > 0:
        #                 p.setPen(pen_2)
        #                 p.drawLine(QtCore.QPointF(i - config.CANDLESTICK_WIDTH, high_price[i]), QtCore.QPointF(i + config.CANDLESTICK_WIDTH, high_price[i]))
                    

        
        p.end()
    
    def paint(self, painter, option, widget):
        """绘制均线"""
        painter.drawPicture(0, 0, self.picture)
    
    def boundingRect(self):
        """返回边界矩形"""
        return QtCore.QRectF(self.picture.boundingRect())


class MovingAverageSlopePeaksManager(BaseChart):
    """移动平均线管理器"""
    
    def __init__(self, plot_widget: pg.PlotWidget):
        """
        初始化移动平均线管理器
        
        Args:
            plot_widget: PyQtGraph绘图控件
        """
        if flags.ENABLE_MOVING_AVERAGE_SLOPE_PEAKS:
            # self.plot_widget = plot_widget
            super().__init__(plot_widget)
            self.ma_items: Dict[int, MovingAverageSlopePeaksItem] = {}
            self.ma_calculator: Optional[MovingAverageSlopePeaksCalculator] = None
            self.data = None
            
            # 默认均线斜率峰值配置
            self.default_periods = [5, 10, 20, 30]
            self.default_colors = {
                5: 'y',      # 5日均线 - 黄色
                10: 'm',     # 10日均线 - 紫色
                20: 'b',     # 20日均线 - 蓝色
                30: 'r'      # 30日均线 - 红色
            }
    
    def set_data(self, data: pd.DataFrame):
        """
        设置数据并初始化计算器
        
        Args:
            data: 股票数据DataFrame
        """
        if not flags.ENABLE_MOVING_AVERAGE_SLOPE_PEAKS:
            return
        self.data = data
        self.ma_calculator = MovingAverageSlopePeaksCalculator(data)
    
    def add_moving_average_slope_peaks(self, periods: List[int] = None, 
                          colors: Dict[int, str] = None,
                          price_column: str = 'close'):
        """
        添加移动平均线斜率峰值
        
        Args:
            periods: 均线周期列表，默认为[5, 10, 20, 30]
            colors: 各周期对应的颜色字典
            price_column: 价格列名
        """
        if not flags.ENABLE_MOVING_AVERAGE_SLOPE_PEAKS:
            return
        if self.ma_calculator is None:
            raise ValueError("请先调用 set_data() 方法设置数据")
        
        # 使用默认配置
        if periods is None:
            periods = self.default_periods
        if colors is None:
            colors = self.default_colors
        
        # 计算均线斜率峰值
        crest_trough_results = self.ma_calculator.calculate_ma_slope_peaks(periods, price_column)
        
        # 移除旧的均线
        self.remove_all_moving_average_slope_peaks()
        
        # 添加新的均线斜率峰值
        for period in periods:
            if period in crest_trough_results:
                print(f"添加均线斜率峰值: {period}")
                crest_trough_data = crest_trough_results[period]
                color = colors.get(period, 'white')
                logger.debug(f"添加均线斜率峰值绘制项颜色: {period}, {color}")
                
                # 创建均线斜率峰值绘制项
                ma_item = MovingAverageSlopePeaksItem(self.data, crest_trough_data, color, self.vb, 5)
                logger.debug(f"添加均线斜率峰值绘制项: {period}, {ma_item}")
                self.ma_items[period] = ma_item
                logger.debug(f"添加均线斜率峰值绘制项: {period}")
                
                # 添加到图表
                self.plot_widget.addItem(ma_item)
                logger.debug(f"添加均线斜率峰值绘制项到图表: {period}")
                
                # 设置层级（在K线下方）
                ma_item.setZValue(50)
    
    def remove_moving_average_slope_peaks(self, period: int):
        """
        移除指定周期的移动平均线
        
        Args:
            period: 均线周期
        """
        if not flags.ENABLE_MOVING_AVERAGE_SLOPE_PEAKS:
            return
        if period in self.ma_items:
            self.plot_widget.removeItem(self.ma_items[period])
            del self.ma_items[period]
    
    def remove_all_moving_average_slope_peaks(self):
        """移除所有移动平均线"""
        if not flags.ENABLE_MOVING_AVERAGE_SLOPE_PEAKS:
            return

        for period in list(self.ma_items.keys()):
            self.remove_moving_average_slope_peaks(period)
    
    def update_moving_average_slope_peaks(self, periods: List[int] = None,
                              colors: Dict[int, str] = None,
                              price_column: str = 'close'):
        """
        更新移动平均线
        
        Args:
            periods: 均线周期列表
            colors: 各周期对应的颜色字典
            price_column: 价格列名
        """
        if not flags.ENABLE_MOVING_AVERAGE_SLOPE_PEAKS:
            return
        self.vb = self.plot_widget.getViewBox()
        self.add_moving_average_slope_peaks(periods, colors, price_column)
    
    def get_ma_slope_peaks_range(self, periods: List[int] = None) -> Tuple[float, float]:
        """
        获取均线斜率峰值的价格范围
        
        Args:
            periods: 均线周期列表
            
        Returns:
            (最小价格, 最大价格)
        """
        if self.ma_calculator is None:
            return 0.0, 0.0
        
        if periods is None:
            periods = list(self.ma_items.keys())
        
        return self.ma_calculator.get_ma_range(periods)
    
    def get_ma_slope_peaks_data(self, period: int) -> Optional[pd.Series]:
        """获取指定周期的均线数据"""
        if self.ma_calculator is None:
            return None
        return self.ma_calculator.get_ma_slope_peaks_data(period)
    
    def get_available_ma_slope_peaks_periods(self) -> List[int]:
        """获取可用的均线周期列表"""
        return list(self.ma_items.keys())
    
    def set_ma_slope_peaks_visible(self, period: int, visible: bool):
        """
        设置指定周期均线的可见性
        
        Args:
            period: 均线周期
            visible: 是否可见
        """
        if not flags.ENABLE_MOVING_AVERAGE_SLOPE_PEAKS:
            return
        if period in self.ma_items:
            self.ma_items[period].setVisible(visible)
    
    def set_all_ma_slope_peaks_visible(self, visible: bool):
        """
        设置所有均线的可见性
        
        Args:
            visible: 是否可见
        """
        if not flags.ENABLE_MOVING_AVERAGE_SLOPE_PEAKS:
            return
        for ma_item in self.ma_items.values():
            ma_item.setVisible(visible)

    def set_x_range(self, x_min: float, x_max: float):
        """
        设置X轴范围
        
        Args:
            x_min: X轴最小值
            x_max: X轴最大值
        """
        if not flags.ENABLE_MOVING_AVERAGE_SLOPE_PEAKS:
            return
        super().set_x_range(x_min, x_max)
    
    def set_date_ticks(self, date_ticks: List[Tuple[int, str]]):
        """
        设置X轴日期刻度
        
        Args:
            date_ticks: 日期刻度列表
        """
        if not flags.ENABLE_MOVING_AVERAGE_SLOPE_PEAKS:
            return
        super().set_date_ticks(date_ticks)

    def set_price_range(self, min_price: float, max_price: float, padding: float = 0.1):
        """
        设置价格范围
        
        Args:
            min_price: 最低价格
            max_price: 最高价格
            padding: 边距比例
        """
        if not flags.ENABLE_MOVING_AVERAGE_SLOPE_PEAKS:
            return
        super().set_price_range(min_price, max_price, padding)